from mcp.server.fastmcp import FastMCP
from sympy import cos, sin
import numpy as np

class RobotKinematics:
    """机器人运动学计算类"""
    def __init__(self, link_lengths=[1.0, 1.0]):
        """
        初始化机器人参数
        :param link_lengths: 各连杆长度列表
        """
        self.link_lengths = link_lengths
        self.mcp = FastMCP("Robot Kinematics Server")

    def dh_matrix(self, theta, d, a, alpha):
        """DH参数矩阵计算"""
        ct = np.cos(theta)
        st = np.sin(theta)
        ca = np.cos(alpha)
        sa = np.sin(alpha)
        return np.array([
            [ct, -st*ca, st*sa, a*ct],
            [st, ct*ca, -ct*sa, a*st],
            [0, sa, ca, d],
            [0, 0, 0, 1]
        ])

# 创建机器人运动学实例
robot = RobotKinematics()
mcp = robot.mcp

@mcp.tool()
def forward_kinematics(joint_angles: list) -> list:
    """
    正运动学求解
    :param joint_angles: 关节角度列表(度)
    :return: 末端执行器位置[x,y,z]
    """
    try:
        # 输入验证
        if len(joint_angles) != 3:
            raise ValueError("关节角度参数必须为3个")
            
        # 将角度转换为弧度
        theta = np.radians(joint_angles)
        
        # 计算变换矩阵链
        T = np.eye(4)
        # 第一关节
        T = T @ robot.dh_matrix(theta=theta[0], d=0, a=0, alpha=0)
        # 第二关节
        T = T @ robot.dh_matrix(theta=theta[1], d=0, a=robot.link_lengths[0], alpha=0)
        # 第三关节
        T = T @ robot.dh_matrix(theta=theta[2], d=0, a=robot.link_lengths[1], alpha=0)
            
        # 返回末端位置
        return T[:3,3].tolist()
    except Exception as e:
        print(f"正运动学计算错误: {str(e)}")
        return []

@mcp.tool()
def inverse_kinematics(target_pose: list) -> list:
    """逆运动学求解"""
    from sympy import symbols, nsolve
    import math
    
    # 创建符号变量
    theta1, theta2, theta3 = symbols('theta1 theta2 theta3')
    x, y, z = target_pose
    
    # 构建正运动学方程
    def forward_eq(theta1, theta2, theta3):
        return [
            math.cos(theta1)*(math.cos(theta2 + theta3) + math.cos(theta2)) - x,
            math.sin(theta1)*(math.cos(theta2 + theta3) + math.cos(theta2)) - y,
            math.sin(theta2 + theta3) + math.sin(theta2) - z
        ]
    
    # 使用数值求解
    try:
        # 改进初始猜测值
        initial_guess = {
            theta1: math.atan2(y, x) if (x != 0 or y != 0) else 0.1,
            theta2: math.pi/4,  # 45度
            theta3: -math.pi/4   # -45度
        }
        
        # 添加约束条件避免奇异解
        solutions = nsolve(
            forward_eq(theta1, theta2, theta3),
            (theta1, theta2, theta3),
            initial_guess,
            dict=True,
            verify=False
        )
        
        if solutions:
            return [float(solutions[0][var]) for var in [theta1, theta2, theta3]]
        return []
    except Exception as e:
        print(f"逆运动学求解失败: {str(e)}")
        return []

def test_kinematics():
    """测试运动学计算"""
    # 测试正运动学
    print("测试正运动学:")
    print(forward_kinematics([30, 45, 0]))
    
    # 测试逆运动学
    print("测试逆运动学:")
    print(inverse_kinematics([1.366, 0, 0]))

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--test", action="store_true", help="运行测试")
    args = parser.parse_args()
    
    if args.test:
        test_kinematics()
    else:
        mcp.run(transport="stdio")